#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# File: qmt_core.py
# Date: 2025/4/17
# Author: zhang shuirong

from walrus import Database
import datetime, time
from walrus import Database
from xtquant import xtdata
from xtquant.xttrader import XtQuantTrader, XtQuantTraderCallback
from xtquant.xttype import StockAccount
from xtquant import xtconstant
import sys, os, json
import logging
import pprint
from decimal import Decimal, ROUND_HALF_UP
import requests
import pandas as pd
from log_config import get_logger

# 获取日志记录器
logger = get_logger()


# 配置参数
class Config:
    DEDUCT_HANDLING_FEE = 0.999914  # 手续费
    QMT_PATH = r'C:\Users\Administrator\sofw\迅投极速策略交易系统交易终端 大同证券QMT实盘\userdata_mini'
    ACCOUNT_ID = '99066915'
    RETRY_TIMES = 3  # 重试次数
    RETRY_INTERVAL = 5  # 重试间隔(秒)
    MAX_ORDER_WAIT_TIME = 120  # 最大订单等待时间(秒)


class MyXtQuantTraderCallback(XtQuantTraderCallback):
    def on_disconnected(self):
        """连接断开"""
        logger.error("连接断开")

    def on_stock_order(self, order):
        """委托回报推送"""
        logger.info(f"订单状态更新: {order.stock_code}, 状态: {order.order_status}, 系统ID: {order.order_sysid}")

    def on_stock_trade(self, trade):
        """成交变动推送"""
        logger.info(f"成交回报: 账户 {trade.account_id}, 代码 {trade.stock_code}, 订单ID {trade.order_id}")

    def on_order_error(self, order_error):
        """委托失败推送"""
        logger.error(
            f"订单失败: 订单ID {order_error.order_id}, 错误码 {order_error.error_id}, 错误信息 {order_error.error_msg}")

    def on_cancel_error(self, cancel_error):
        """撤单失败推送"""
        logger.error(
            f"撤单失败: 订单ID {cancel_error.order_id}, 错误码 {cancel_error.error_id}, 错误信息 {cancel_error.error_msg}")

    def on_order_stock_async_response(self, response):
        """异步下单回报推送"""
        logger.info(f"异步下单响应: 账户 {response.account_id}, 订单ID {response.order_id}, 序列号 {response.seq}")

    def on_account_status(self, status):
        """账户状态更新"""
        logger.info(f"账户状态: 账户ID {status.account_id}, 类型 {status.account_type}, 状态 {status.status}")


def setup_trader():
    """初始化交易连接"""
    xt_trader = XtQuantTrader(Config.QMT_PATH, int(datetime.datetime.now().strftime("%Y%m%d%H%M%S")))
    callback = MyXtQuantTraderCallback()
    xt_trader.register_callback(callback)
    xt_trader.start()

    if not connect_to_trader(xt_trader):
        logger.error("无法连接QMT交易系统")
        sys.exit(1)

    return xt_trader, StockAccount(Config.ACCOUNT_ID)


def connect_to_trader(xt_trader):
    """连接交易系统，带重试机制"""
    for attempt in range(1, Config.RETRY_TIMES + 1):
        if xt_trader.connect() == 0:
            logger.info("成功连接miniQMT")
            return True
        logger.warning(f"miniQMT连接失败，重试 {attempt}/{Config.RETRY_TIMES}")
        time.sleep(Config.RETRY_INTERVAL * attempt)
    return False


def get_position_dicts(positions):
    """获取持仓字典"""
    position_total_dict = {i.stock_code: i.m_nVolume for i in positions}
    position_available_dict = {i.stock_code: i.m_nCanUseVolume for i in positions}
    logger.info("持仓字典: %s", pprint.pformat(position_total_dict))
    logger.info("可用持仓字典: %s", pprint.pformat(position_available_dict))
    return position_total_dict, position_available_dict


def determine_trade_actions(stock_list, position_available_dict, exclude_codes_list):
    """确定交易操作"""
    position_available_dict = {key: value for key, value in position_available_dict.items() if value != 0}
    avai_stock_set = set(position_available_dict.keys())
    buy_code_list = [code for code in stock_list if code not in avai_stock_set]
    sell_code_list = [code for code in avai_stock_set if code not in stock_list and code.startswith('1')]
    sell_code_list = [code for code in sell_code_list if code not in exclude_codes_list]
    continue_code_list = [code for code in stock_list if code in avai_stock_set] + exclude_codes_list
    logger.info("今日买入标的：%s", buy_code_list)
    logger.info("今日卖出标的：%s", sell_code_list)
    logger.info("今日继续持有标的：%s", continue_code_list)
    return buy_code_list, sell_code_list, continue_code_list


def execute_sell_orders(xt_trader, acc, sell_code_list, position_available_dict):
    """执行卖出订单"""
    fix_result_order_id_list = []  # 存储卖出订单ID的列表
    for stock in sell_code_list:  # 遍历需要卖出的股票代码列表
        available_vol = position_available_dict.get(stock, 0)  # 获取该股票的可用持仓数量
        if available_vol > 0:  # 如果可用数量大于0，则执行卖出操作
            sell_price = xtdata.get_full_tick([stock])[stock]['bidPrice'][2]  # 获取卖二价格
            order_id = xt_trader.order_stock(acc, stock, xtconstant.STOCK_SELL, available_vol, xtconstant.FIX_PRICE,
                                             sell_price, 'runshit', '证券卖出')  # 下单卖出
            fix_result_order_id_list.append(order_id)  # 将订单ID添加到列表
            logger.info("卖出订单: %s", order_id)  # 记录卖出订单信息
        else:
            logger.info("无法卖出 %s , 可用数量为0", stock)  # 如果可用数量为0，记录日志
    return fix_result_order_id_list  # 返回卖出订单ID列表


def wait_for_orders_to_complete(xt_trader, acc, order_id_list):
    """等待订单完成"""
    start_time = time.time()
    while time.time() - start_time < 120:
        try:
            if all(xt_trader.query_stock_order(acc, order_id).order_status == 56 for order_id in order_id_list):
                logger.info("全部卖出成功")
                return
            time.sleep(1)
        except:
            logger.info('捕获到异常，不过问题不大，直接跳过')
            time.sleep(1)
    logger.warning("120秒内卖出失败")


def execute_buy_orders(xt_trader, acc, buy_dict, total_asset):
    """执行买入订单（按仓位比例分配资金）"""
    buy_price = {}
    buy_vol = {}

    # 计算总权重（仓位比例总和）
    total_weight = sum(buy_dict.values())
    if total_weight <= 0:
        logger.error("仓位权重总和必须大于0")
        return {}, {}

    # 计算总可用资金（扣除手续费）
    total_cash = total_asset * Config.DEDUCT_HANDLING_FEE

    # 按权重分配资金
    for stock, weight in buy_dict.items():
        # 计算该股票应分配金额
        stock_cash = total_cash * weight

        # 获取买一价或最新价
        logger.info(f"{stock} 全推行情： {xtdata.get_full_tick([stock])}")
        buy_price[stock] = xtdata.get_full_tick([stock])[stock]['askPrice'][1] or xtdata.get_full_tick([stock])[
            'lastPrice']

        # 计算买入数量（按10股整数倍）
        if buy_price[stock] > 0:
            # 修改后的买入数量计算（修正为100整数倍）
            buy_quantity = max(100, int(stock_cash / buy_price[stock] // 100) * 100)
        else:
            logger.error(f"{stock} 价格异常: {buy_price[stock]}")
            continue

        # 执行买入订单
        order_id = xt_trader.order_stock(
            acc,
            stock,
            xtconstant.STOCK_BUY,
            buy_quantity,
            xtconstant.FIX_PRICE,
            buy_price[stock],
            'runshit',
            '证券买入'
        )

        # 记录交易信息
        buy_vol[stock] = buy_quantity
        logger.info(
            f"买入订单[{order_id}]: {stock} "
            f"价格={buy_price[stock]:.2f} "
            f"数量={buy_quantity} "
            f"金额={buy_price[stock] * buy_quantity:.2f}"
        )

    return buy_price, buy_vol


def qmt_execute(jukuan_positions, exclude_codes_list):
    logger.info("今日交易开始")
    logger.info("处理聚宽数据--")
    stock_list, modified_stock_dict = jukuan_positions_handle(jukuan_positions)
    try:
        xt_trader, acc = setup_trader()
        positions = xt_trader.query_stock_positions(acc)
        position_total_dict, position_available_dict = get_position_dicts(positions)

        buy_code_list, sell_code_list, continue_code_list = determine_trade_actions(
            stock_list, position_available_dict, exclude_codes_list
        )
        logger.info(f"最终买入的 buy_code_list: {buy_code_list}")
        logger.info(f"最终卖出的 sell_code_list: {sell_code_list}")

        if sell_code_list:
            order_ids = execute_sell_orders(xt_trader, acc, sell_code_list, position_available_dict)
            wait_for_orders_to_complete(xt_trader, acc, order_ids)

        buy_dict = {key: modified_stock_dict[key] for key in buy_code_list}
        logger.info(f"按照仓位比例买入字典: {buy_dict}")
        total_asset = xt_trader.query_stock_asset(acc).total_asset
        if buy_code_list:
            execute_buy_orders(xt_trader, acc, buy_dict, total_asset)

        logger.info("今日交易结束")

    except Exception as e:
        logger.error(f"交易执行失败: {str(e)}")
        sys.exit(1)


def jukuan_positions_handle(jukuan_positions):
    stock_dict = jukuan_positions
    stock_list = list(stock_dict.keys())
    # 修改后缀
    stock_list = [symbol.replace('XSHG', 'SH').replace('XSHE', 'SZ') for symbol in stock_list]
    logger.info(f"提取的 symbol 列表: {stock_list}")

    # 创建新的字典，键为修改后的后缀
    modified_stock_dict = {symbol.replace('XSHG', 'SH').replace('XSHE', 'SZ'): quantity for symbol, quantity in
                           stock_dict.items()}
    logger.info(f"修改后的 stock_info: {modified_stock_dict}")
    return stock_list, modified_stock_dict


if __name__ == '__main__':
    pass
